import asyncio
import json
import websockets
import random
import time
import queue

# WebSocket模块：处理与服务器的通信

# 全局标志：用户是否主动退出
user_initiated_exit = False

class WebSocketManager:
    def __init__(self):
        self.websocket_url = "wss://api.tenclass.net/xiaozhi/v1/"
        self.token = "test-token"
        self.session_id = ''
        self.MAX_RECONNECT_ATTEMPTS = 5
        self.connection_established = False
        self.headers = {
            "Authorization": f"Bearer {self.token}",
            "Protocol-Version": "1",
            "Device-Id": "94:a9:90:1a:e2:08",
            "Client-Id": "device-uuid-1234567"
        }
    
    # 心跳任务：定期发送ping保持连接活跃
    async def heartbeat(self, websocket):
        try:
            while True:
                await asyncio.sleep(20)  # 20秒发送一次心跳
                try:
                    # 使用send_ping方法而非自定义ping消息
                    await websocket.ping()
                except Exception as ping_e:
                    print(f"[WARN] 发送心跳包失败: {ping_e}")
                    break
        except asyncio.CancelledError:
            pass
    
    # 主连接逻辑
    async def connect_websocket(self, audio_queue_callback, recording_queue_callback, return_websocket=False):
        global user_initiated_exit
        user_initiated_exit = False
        reconnect_attempts = 0
        
        if return_websocket:
            # 直接返回websocket实例给调用者处理
            try:
                websocket = await websockets.connect(
                    self.websocket_url, 
                    extra_headers=self.headers,
                    ping_interval=20,
                    ping_timeout=15,
                    close_timeout=5,
                    open_timeout=20,
                    max_size=2**20,
                    compression=None
                )
                print("连接成功！")
                self.connection_established = True
                
                # 发送hello消息
                hello_msg = {
                    "type": "hello",
                    "version": 1,
                    "features": {"mcp": False},
                    "transport": "websocket",
                    "audio_params": {
                        "format": "opus",
                        "sample_rate": 16000,
                        "channels": 1,
                        "frame_duration": 60
                    }
                }
                await websocket.send(json.dumps(hello_msg))
                print("已发送hello消息")
                
                # 等待hello应答并获取session_id
                try:
                    response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    if isinstance(response, str):
                        try:
                            response_json = json.loads(response)
                            print(f"收到hello应答: {response_json}")
                            self.session_id = response_json.get('session_id')
                        except json.JSONDecodeError:
                            print("[ERROR] hello应答无法解析为JSON")
                    else:
                        print("[ERROR] hello应答不是字符串格式")
                except asyncio.TimeoutError:
                    print("[WARN] 超时未收到hello应答，继续运行")
                
                return websocket
            except Exception as e:
                print(f"[ERROR] 连接出错: {e}")
                self.connection_established = False
                return None
        
        # 原有模式：在方法内部处理连接生命周期和重连逻辑
        while True:
            try:
                # 优化WebSocket连接参数
                async with websockets.connect(
                    self.websocket_url, 
                    extra_headers=self.headers,
                    ping_interval=20,
                    ping_timeout=15,
                    close_timeout=5,
                    open_timeout=20,
                    max_size=2**20,
                    compression=None
                ) as websocket:
                    print("连接成功！")
                    self.connection_established = True
                    reconnect_attempts = 0
                    
                    # 发送hello消息
                    hello_msg = {
                        "type": "hello",
                        "version": 1,
                        "features": {"mcp": False},
                        "transport": "websocket",
                        "audio_params": {
                            "format": "opus",
                            "sample_rate": 16000,
                            "channels": 1,
                            "frame_duration": 60
                        }
                    }
                    await websocket.send(json.dumps(hello_msg))
                    print("已发送hello消息")
                    
                    # 等待hello应答并获取session_id
                    try:
                        response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                        if isinstance(response, str):
                            try:
                                response_json = json.loads(response)
                                print(f"收到hello应答: {response_json}")
                                self.session_id = response_json.get('session_id')
                            except json.JSONDecodeError:
                                print("[ERROR] hello应答无法解析为JSON")
                        else:
                            print("[ERROR] hello应答不是字符串格式")
                    except asyncio.TimeoutError:
                        print("[WARN] 超时未收到hello应答，继续运行")
                    
                    # 创建并返回任务列表
                    tasks = await self.create_tasks(websocket, audio_queue_callback, recording_queue_callback)
                    
                    # 等待任一任务完成
                    done, pending = await asyncio.wait(
                        tasks,
                        return_when=asyncio.FIRST_COMPLETED,
                        timeout=None
                    )
                    
                    # 检查完成的任务是否有异常
                    error_occurred = False
                    for task in done:
                        if task.get_name() != "message_sender" and task.exception():
                            error_type = type(task.exception()).__name__
                            print(f"[ERROR] 任务 {task.get_name()} 出现异常: {task.exception()}")
                            print(f"[ERROR] 异常类型: {error_type}")
                            if "ConnectionClosed" in error_type:
                                print("[WARN] WebSocket连接已关闭，将尝试重连")
                            error_occurred = True
                    
                    # 取消所有剩余的任务
                    for task in pending:
                        task.cancel()
                        try:
                            await task
                        except asyncio.CancelledError:
                            print(f"任务 {task.get_name()} 已取消")
                    
                    # 只有当用户输入q主动退出时才不重连
                    if not error_occurred:
                        print("[INFO] 检测到可能是任务完成导致的关闭，准备重新连接以保持对话状态...")
            
            except websockets.exceptions.ConnectionClosed as e:
                print(f"[ERROR] WebSocket连接已关闭: {e}")
                if "1005" in str(e):
                    print("[WARN] 检测到1005错误（内部错误），将尝试重连")
                reconnect_attempts += 1
            except websockets.exceptions.ConnectionClosedError as e:
                print(f"[ERROR] WebSocket连接错误: {e}")
                reconnect_attempts += 1
            except websockets.exceptions.ConnectionClosedOK as e:
                print(f"[INFO] WebSocket连接已正常关闭: {e}")
                print("[INFO] 准备重新连接以保持对话状态...")
                reconnect_attempts += 1
            except Exception as e:
                error_type = type(e).__name__
                print(f"[ERROR] 连接出错: {e}")
                print(f"[ERROR] 错误类型: {error_type}")
                reconnect_attempts += 1
            finally:
                self.connection_established = False
            
            # 检查用户是否主动退出
            if user_initiated_exit:
                print("[INFO] 用户主动退出，停止重连")
                break
                
            # 检查是否达到最大重连次数
            if reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS:
                print(f"[ERROR] 已达到最大重连次数({self.MAX_RECONNECT_ATTEMPTS})，停止重连")
                break
            
            # 智能退避重连
            base_delay = min(15, 2 ** (reconnect_attempts - 1))
            jitter = random.uniform(0.5, 1.5)
            reconnect_delay = base_delay * jitter
            print(f"[INFO] {reconnect_delay:.1f}秒后尝试第{reconnect_attempts}次重连...")
            await asyncio.sleep(reconnect_delay)
        
        return self.session_id
    
    # 创建任务列表（这个方法将在主模块中实现，以便整合所有功能）
    async def create_tasks(self, websocket, audio_queue_callback, recording_queue_callback):
        # 这个方法将由主模块覆盖
        pass

# 消息监听函数（将在主模块中与音频处理结合）
async def listen_messages(websocket, audio_queue, is_first_audio_packet):
    audio_packet_count = 0
    consecutive_audio_blocks = 0
    last_audio_time = time.time()
    
    try:
        while True:
            current_time = time.time()
            
            # 接收消息（带超时）
            try:
                message = await asyncio.wait_for(websocket.recv(), timeout=30.0)
            except asyncio.TimeoutError:
                print("⏱️  接收消息超时，继续等待...")
                continue
            
            # 处理不同类型的消息
            if isinstance(message, str):
                # 文本消息处理
                try:
                    data = json.loads(message)
                    # 这里可以添加更多消息类型的处理
                    print(f"收到文本消息: {data}")
                except json.JSONDecodeError:
                    print(f"[ERROR] 无法解析JSON消息: {message}")
            else:
                # 二进制消息处理（假设是Opus音频数据）
                audio_packet_count += 1
                last_audio_time = current_time
                consecutive_audio_blocks += 1
                
                # 将音频数据放入队列
                try:
                    # 尝试使用异步方式添加到队列
                    await asyncio.wait_for(audio_queue.put(message), timeout=0.1)
                except asyncio.QueueFull:
                    print("[ERROR] 音频队列已满，丢弃部分数据")
                    # 尝试清理一个位置并添加新数据
                    try:
                        # 尝试异步获取
                        await asyncio.wait_for(audio_queue.get(), timeout=0.01)
                        audio_queue.task_done()
                        # 异步添加
                        await asyncio.wait_for(audio_queue.put(message), timeout=0.1)
                        print("[OK] 已清理一个位置并添加新数据")
                    except (asyncio.TimeoutError, asyncio.QueueEmpty, asyncio.QueueFull):
                        print("[ERROR] 处理音频数据时出现队列错误")
                    except Exception as put_e:
                        print(f"[ERROR] 将音频数据加入队列失败: {put_e}")
    
    except websockets.exceptions.ConnectionClosed as e:
        print(f"🔌 连接已关闭，停止监听 - 原因: {e}")
    except Exception as e:
        print(f"❌ 监听出错: {e}")
        import traceback
        traceback.print_exc()
    finally:
        print(f"📋 监听任务结束 - 共接收 {audio_packet_count} 个音频包")